__init__.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177
  1. # Copyright (c) 2012 Giorgos Verigakis <verigak@gmail.com>
  2. #
  3. # Permission to use, copy, modify, and distribute this software for any
  4. # purpose with or without fee is hereby granted, provided that the above
  5. # copyright notice and this permission notice appear in all copies.
  6. #
  7. # THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
  8. # WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
  9. # MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
  10. # ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
  11. # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
  12. # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
  13. # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
  14. from __future__ import division, print_function
  15. from collections import deque
  16. from datetime import timedelta
  17. from math import ceil
  18. from sys import stderr
  19. try:
  20. from time import monotonic
  21. except ImportError:
  22. from time import time as monotonic
  23. __version__ = '1.5'
  24. HIDE_CURSOR = '\x1b[?25l'
  25. SHOW_CURSOR = '\x1b[?25h'
  26. class Infinite(object):
  27. file = stderr
  28. sma_window = 10 # Simple Moving Average window
  29. check_tty = True
  30. hide_cursor = True
  31. def __init__(self, message='', **kwargs):
  32. self.index = 0
  33. self.start_ts = monotonic()
  34. self.avg = 0
  35. self._avg_update_ts = self.start_ts
  36. self._ts = self.start_ts
  37. self._xput = deque(maxlen=self.sma_window)
  38. for key, val in kwargs.items():
  39. setattr(self, key, val)
  40. self._width = 0
  41. self.message = message
  42. if self.file and self.is_tty():
  43. if self.hide_cursor:
  44. print(HIDE_CURSOR, end='', file=self.file)
  45. print(self.message, end='', file=self.file)
  46. self.file.flush()
  47. def __getitem__(self, key):
  48. if key.startswith('_'):
  49. return None
  50. return getattr(self, key, None)
  51. @property
  52. def elapsed(self):
  53. return int(monotonic() - self.start_ts)
  54. @property
  55. def elapsed_td(self):
  56. return timedelta(seconds=self.elapsed)
  57. def update_avg(self, n, dt):
  58. if n > 0:
  59. xput_len = len(self._xput)
  60. self._xput.append(dt / n)
  61. now = monotonic()
  62. # update when we're still filling _xput, then after every second
  63. if (xput_len < self.sma_window or
  64. now - self._avg_update_ts > 1):
  65. self.avg = sum(self._xput) / len(self._xput)
  66. self._avg_update_ts = now
  67. def update(self):
  68. pass
  69. def start(self):
  70. pass
  71. def clearln(self):
  72. if self.file and self.is_tty():
  73. print('\r\x1b[K', end='', file=self.file)
  74. def write(self, s):
  75. if self.file and self.is_tty():
  76. line = self.message + s.ljust(self._width)
  77. print('\r' + line, end='', file=self.file)
  78. self._width = max(self._width, len(s))
  79. self.file.flush()
  80. def writeln(self, line):
  81. if self.file and self.is_tty():
  82. self.clearln()
  83. print(line, end='', file=self.file)
  84. self.file.flush()
  85. def finish(self):
  86. if self.file and self.is_tty():
  87. print(file=self.file)
  88. if self.hide_cursor:
  89. print(SHOW_CURSOR, end='', file=self.file)
  90. def is_tty(self):
  91. return self.file.isatty() if self.check_tty else True
  92. def next(self, n=1):
  93. now = monotonic()
  94. dt = now - self._ts
  95. self.update_avg(n, dt)
  96. self._ts = now
  97. self.index = self.index + n
  98. self.update()
  99. def iter(self, it):
  100. with self:
  101. for x in it:
  102. yield x
  103. self.next()
  104. def __enter__(self):
  105. self.start()
  106. return self
  107. def __exit__(self, exc_type, exc_val, exc_tb):
  108. self.finish()
  109. class Progress(Infinite):
  110. def __init__(self, *args, **kwargs):
  111. super(Progress, self).__init__(*args, **kwargs)
  112. self.max = kwargs.get('max', 100)
  113. @property
  114. def eta(self):
  115. return int(ceil(self.avg * self.remaining))
  116. @property
  117. def eta_td(self):
  118. return timedelta(seconds=self.eta)
  119. @property
  120. def percent(self):
  121. return self.progress * 100
  122. @property
  123. def progress(self):
  124. return min(1, self.index / self.max)
  125. @property
  126. def remaining(self):
  127. return max(self.max - self.index, 0)
  128. def start(self):
  129. self.update()
  130. def goto(self, index):
  131. incr = index - self.index
  132. self.next(incr)
  133. def iter(self, it):
  134. try:
  135. self.max = len(it)
  136. except TypeError:
  137. pass
  138. with self:
  139. for x in it:
  140. yield x
  141. self.next()